import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime
from sqlalchemy import create_engine, text
import time
from functools import wraps
from db_config import DBConfig

# -------------------- 公共方法 --------------------

def retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,)):
    """
    简单的重试装饰器
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            attempts = 0
            current_delay = delay
            while attempts < max_attempts:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    attempts += 1
                    if attempts == max_attempts:
                        print(f"函数 {func.__name__} 重试 {max_attempts} 次后仍然失败: {e}")
                        raise
                    print(f"函数 {func.__name__} 第 {attempts} 次失败: {e}, {current_delay}秒后重试...")
                    time.sleep(current_delay)
                    current_delay *= backoff
        return wrapper
    return decorator

def clean_stock_code(stock_code: str):
    """
    清理股票代码，移除.sz和.sh后缀
    """
    return stock_code.replace('.SZ', '').replace('.SH', '').replace('.BJ', '')

# -------------------- 获取 & 处理数据 --------------------

@retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,))
def get_stock_dividend_data(symbol: str) -> pd.DataFrame:
    """
    获取分红送转数据（带重试）
    """
    try:
        df = ak.stock_fhps_detail_em(symbol=symbol)
        return df
    except Exception as e:
        print(f"获取 {symbol} 分红送转数据异常: {e}")
        return pd.DataFrame()

def prepare_stock_dividend_data(raw_df: pd.DataFrame, stock_code: str, stock_name: str) -> pd.DataFrame:
    """
    清洗分红送转数据，映射到 finance_dividend 表字段
    """
    if raw_df.empty:
        return pd.DataFrame()

    column_mapping = {
        "报告期": "report_date",
        "业绩披露日期": "announce_date",
        "送转股份-送转总比例": "bonus_share_ratio",
        "送转股份-送股比例": "bonus_share_ratio_send",
        "送转股份-转股比例": "bonus_share_ratio_transfer",
        "现金分红-现金分红比例": "cash_dividend_ratio",
        "现金分红-现金分红比例描述": "cash_dividend_desc",
        "现金分红-股息率": "dividend_yield",
        "每股收益": "eps",
        "每股净资产": "bvps",
        "每股公积金": "capital_reserve_per_share",
        "每股未分配利润": "undistributed_profit_per_share",
        "净利润同比增长": "net_profit_growth_rate",
        "总股本": "total_share_capital",
        "预案公告日": "plan_announce_date",
        "股权登记日": "register_date",
        "除权除息日": "ex_dividend_date",
        "方案进度": "plan_progress",
        "最新公告日期": "latest_announce_date"
    }

    df = raw_df.rename(columns=column_mapping)

    # 补充股票代码和名称
    df["stock_code"] = stock_code
    df["stock_name"] = stock_name

    # 日期字段转换
    date_cols = ["report_date", "announce_date", "plan_announce_date",
                 "register_date", "ex_dividend_date", "latest_announce_date"]
    for col in date_cols:
        if col in df.columns:
            df[col] = pd.to_datetime(df[col], errors="coerce").dt.date

    # 数值字段处理
    numeric_cols = ["bonus_share_ratio", "bonus_share_ratio_send", "bonus_share_ratio_transfer",
                    "cash_dividend_ratio", "dividend_yield", "eps", "bvps",
                    "capital_reserve_per_share", "undistributed_profit_per_share",
                    "net_profit_growth_rate", "total_share_capital"]
    for col in numeric_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
            df[col] = df[col].replace([np.inf, -np.inf], np.nan)

    return df

# -------------------- 数据库写入 --------------------

@retry(max_attempts=3, delay=1, backoff=2, exceptions=(Exception,))
def insert_stock_dividend_to_mysql(df: pd.DataFrame) -> int:
    """
    插入或更新 MySQL - finance_dividend（幂等，唯一键：stock_code+register_date）
    """
    if df.empty:
        return 0

    db_config = DBConfig.config
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )

    try:
        with engine.begin() as conn:
            for _, row in df.iterrows():
                row_dict = {k: (None if pd.isna(v) else v) for k, v in row.to_dict().items()}

                columns = ", ".join(row_dict.keys())
                placeholders = ", ".join([f":{k}" for k in row_dict.keys()])

                update_clause = ", ".join([f"{col}=VALUES({col})"
                                           for col in row_dict.keys()
                                           if col not in ("stock_code", "register_date")])

                sql = f"""
                    INSERT INTO finance_dividend ({columns})
                    VALUES ({placeholders})
                    ON DUPLICATE KEY UPDATE {update_clause}
                """
                conn.execute(text(sql), row_dict)

        return len(df)
    except Exception as e:
        print(f"插入/更新分红送转数据失败: {e}")
        raise
    finally:
        engine.dispose()

# -------------------- 同步控制函数 --------------------

def sync_stock_dividend_to_db(stock_code: str, stock_name: str):
    """
    获取 + 清洗 + 入库
    """
    try:
        clean_code = clean_stock_code(stock_code)

        raw_df = get_stock_dividend_data(stock_code)
        if raw_df.empty:
            return {"success": True, "message": f"{stock_name}({clean_code}) 未获取到分红送转数据", "records": 0}

        df = prepare_stock_dividend_data(raw_df, clean_code, stock_name)
        if df.empty:
            return {"success": True, "message": f"{stock_name}({clean_code}) 数据清洗后为空", "records": 0}

        records = insert_stock_dividend_to_mysql(df)
        return {"success": True, "message": f"成功同步 {stock_name}({clean_code}) 分红送转数据 {records} 条", "records": records}

    except Exception as e:
        return {"success": False, "message": f"分红送转数据处理失败: {e}", "records": 0}

def batch_sync_stock_dividend(stock_list):
    """
    批量同步分红送转数据
    """
    total_records = 0
    success_count = 0
    fail_count = 0

    for code, name in stock_list:
        try:
            result = sync_stock_dividend_to_db(code, name)
            if result["success"]:
                success_count += 1
                total_records += result["records"]
                if result["records"] > 0:
                    print(result["message"])
                else:
                    print(f"{name}({clean_stock_code(code)}) 无新数据可同步")
            else:
                fail_count += 1
                print(f"同步 {name}({clean_stock_code(code)}) 失败: {result['message']}")
        except Exception as e:
            fail_count += 1
            print(f"同步 {name}({clean_stock_code(code)}) 异常: {e}")

    print("\n同步完成统计:")
    print(f"成功同步股票数: {success_count}")
    print(f"失败股票数: {fail_count}")
    print(f"总记录数: {total_records}")

# -------------------- 使用示例 --------------------
if __name__ == "__main__":
    stock_list = [
        ("002236", "大华股份"),
        # ("600000.SH", "浦发银行"),
    ]
    batch_sync_stock_dividend(stock_list)
    print("分红送转数据同步完成")
